import random
import time
import re
import paramiko
import threading
from queue import Queue
from io import BytesIO
from itertools import count

RETRY_TIMES = 3
RECV_BUFFER_SIZE = 4096

MESSAGE_LOGIN = 1
MESSAGE_LOGOUT = 2
MESSAGE_COMMAND = 3


def remove_ansi_escape_sequences(text: str):
    """
    去除字符串中的 ANSI 转义序列
    """
    s = text.rsplit('\x1b]', 1)[0]
    pattern = r'\x1B[@-_][0-?]*[ -/]*[@-~]'
    # \x1B \x1b匹配字符串中的ESC字符即\x1b
    # [@-_] 从ASCII表中从@到_的所有字符出现一次，即@A-Z[\]^_，不包含数字和小写
    # [0-?]* 从ASCII表中从0到?的所有字符出现任意次，即0-9:;<=>?，不包含大小写
    # [ -/]* 从ASCII表中从空格到/所有字符出现任意次，即空格!"#$%&'()*+,-./，不包含字母和数字
    # [@-~] 从ASCII表中从@到~的所有字符出现一次，包含大小写，不包含数字
    regex = re.compile(pattern)
    return regex.sub('', s)

def strip_ansi_bytes(b: bytes) -> bytes:
    """去除 bytes 中的 ANSI 转义序列"""
    ansi_escape = re.compile(rb'\x1B[@-_][0-?]*[ -/]*[@-~]')
    return ansi_escape.sub(b'', b)

class Message(dict):
    def __init__(self, typ: int, command: str = "", output: str = "", ts: float = None):
        """typ 0welcome；1exit；2command;"""
        super().__init__()
        self["type"] = typ
        self["command"] = command
        self["output"] = output
        self["created"] = ts or time.time()


class SSHClient:
    """
    status_switch用来判断是命令输入模式，还是交互模式，直到接收到命令提示符交互模式结束又进入命令模式
    _cmdq用来缓冲命令及其输出，_command方法用来处理该队列，把合法数据凑成命令类型Message
    """

    def __init__(self, hostname, username, port=22, password='', ssh_pkey_path=''):
        self._client = paramiko.SSHClient()
        self._conn_kwargs = {
            "hostname": hostname,
            "username": username,
            "port": port,
        }
        if password:
            self._conn_kwargs["password"] = password
        if ssh_pkey_path:
            self._conn_kwargs["pkey"] = paramiko.RSAKey.from_private_key_file(ssh_pkey_path)
        if not password and not ssh_pkey_path:
            raise Exception("连接参数异常，密码和密钥至少提供一个")

        # 初始化队列，来异步处理命令、接收数据
        self._sendq = Queue()  # 用户输入命令队列
        self._recvq = Queue()  # 远程主机输出队列（返回给前端）
        self._cmdq = Queue()  # 命令及其结果的队列
        self._msgq = Queue()  # 输出审计消息的队列
        self.channel = self._initialize()  # 初始化连接，启动线程
        self._closed = False

        self.status_switch = True  # True命令模式，输出是命令；False交互模式，输入都不是命令
        self._cmd_id = 0  # 命令id
        self._cmd_id_counter = count(random.randint(1, 1000))  # 命令id生成器

        threading.Thread(target=self._send).start()
        threading.Thread(target=self._recv).start()
        threading.Thread(target=self._command).start()
        threading.Thread(target=self._watch).start()  # 监视
        # 创建成功，可以发一条登录成功的消息
        msg = Message(MESSAGE_LOGIN)
        self._msgq.put(msg)

    def _initialize(self) -> paramiko.Channel:
        """连接，开启交互shell，非阻塞，返回channel"""
        self._client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        self._client.connect(**self._conn_kwargs)
        channel = self._client.invoke_shell("xterm")
        channel.setblocking(False)  # 非阻塞模式，如果没有数据，避免 _recv 线程卡住
        return channel

    def get_next_cmd_id(self):
        """从计数器 _cmd_id_counter 中获取下一个命令 ID，用于唯一标识每个命令"""
        self._cmd_id = next(self._cmd_id_counter)
        return self._cmd_id

    @property
    def cmd_id(self):
        """属性方法，返回当前命令 ID"""
        return self._cmd_id

    def send(self, chars: str):
        """接收用户的命令，将用户输入（字符串或 bytes）放入 _sendq 队列，等待 _send 线程处理"""
        if isinstance(chars, str):
            self._sendq.put(chars.encode())
        elif isinstance(chars, bytes):
            self._sendq.put(chars)
        else:
            pass  # 如有所需，以后扩展

    def _send(self):
        """
        监听 _sendq，获取用户输入。
        在“命令模式”下累积输入，直到遇到 \r（回车），表示一条命令完成。
        将完整命令记录进 _cmdq，并切换为“交互模式”。
        实际通过 channel.send() 把每个按键发给远程服务器（模拟键盘输入）。
        """
        input_ = BytesIO()  # 记录按键，遇到一个回车\r，算一个命令
        try:
            # 开始发送命令和接收该命令输入结果
            while not self.channel.closed:
                if self._sendq.empty():
                    time.sleep(0.1)
                    continue
                key: bytes = self._sendq.get()

                if self.status_switch == True:  # 命令模式
                    input_.write(key)  # 凑命令
                    if key.endswith(b'\r'):  # 命令模式下的\r之前算是一个命令
                        cmd_id = self.get_next_cmd_id()  # 获取一个新的命令id
                        self.status_switch = False  # status_switch再次为True之前，所有输入都认为是交互而不是命令
                        cmd = input_.getvalue().strip().decode()
                        self._cmdq.put((  # 命令打入队列
                            cmd_id, 1,  # 类型1表示命令
                            cmd, time.time()  # 时间戳
                        ))
                        input_.seek(0)
                        input_.truncate()  # 清空
                self.channel.send(key)  # 收到的所有输入都发到SSH服务端
        except Exception as e:
            print(e)
        finally:
            self.close()
            print("send thread over")

    def _recv(self):
        """
        所有原始数据放入 _recvq
        在“交互模式”下，将输出内容暂存
        当检测到 $ 或 # 提示符时，认为命令执行完成，将输出打包放入 _cmdq。
        """
        output_ = BytesIO()  # 记录命令结果
        try:
            while not self.channel.closed:
                if not self.channel.recv_ready():  # channel关闭，recv_ready返回false，不报错
                    time.sleep(0.1)
                    continue
                data: bytes = self.channel.recv(RECV_BUFFER_SIZE)  # 非阻塞到有数据到来

                # 一条命令的返回的结果，可能需要多次从缓冲区拿回
                # data有3种：1、欢迎信息 2、用户输入的键的回显 3、命令的结果
                self._recvq.put(data)  # 此队列里面的信息直接返回给浏览器端

                if self.status_switch == False:  # 进入交互模式，用户输入的键都是输出而不是命令
                    output_.write(data)

                # invoke_shell下只能从返回的内容判断是否是命令提示行
                clean_data = strip_ansi_bytes(data).rstrip()
                if clean_data.endswith((b'$', b'#')):  # 如果遇到命令提示符一定是执行完了
                    cmd_id = self.cmd_id  # 获取当前命令id
                    self.status_switch = True  # 切回命令模式，用户的输入就是命令
                    output = output_.getvalue().decode()
                    self._cmdq.put((
                        cmd_id, 2,  # 类型2表示命令的输出
                        output, time.time()
                    ))
                    output_.seek(0)
                    output_.truncate()  # 清空
        except Exception as e:
            print(e)
        finally:
            self.close()
            print("recv thread over")

    def output(self):
        """所有命令的回显及其输出"""
        while not self.closed:
            if not self._recvq.empty():
                return self._recvq.get()  # 能拿的时候已经不阻塞了 bytes
            else:
                time.sleep(0.1)

    def message(self):
        # 返回各种审计消息
        while True:
            if not self._msgq.empty():
                msg = self._msgq.get()  # 能拿的时候已经不阻塞了 msg
                self._msgq.task_done()  # 让unfinished_tasks减1
                return msg
            elif self._closed:
                return None
            else:
                time.sleep(0.1)

    def _command(self):
        """从 _cmdq 中按顺序取出“命令”和“输出”，匹配 ID 后合并为 Message"""
        # ssh的命令和结果，所以这个整理消息的事SSHClient做
        # 如果出现前面是105 1后一个号是106 1，那就，前面的就丢弃105
        while not self.closed:
            (cid, typ, cmd, ts) = self._cmdq.get()  # 能拿的时候已经不阻塞了 msg
            if typ != 1:  # 命令
                continue
            (rid, typ, output, _) = self._cmdq.get()
            if typ != 2:  # 输出
                continue
            if cid != rid:
                continue
            # 组成一个消息
            msg = Message(MESSAGE_COMMAND, cmd, remove_ansi_escape_sequences(output))
            self._msgq.put(msg)
            print(msg)
            print('~' * 30)

    def _watch(self):
        """监视channel是否关闭，来关闭当前对象"""
        while not self.channel.closed:
            time.sleep(0.5)
        self.close()

    @property
    def closed(self):
        return self._closed

    def close(self):
        if self.closed:
            return
        # 退出消息
        msg = Message(MESSAGE_LOGOUT)
        self._msgq.put(msg)  # 每增加一个元素，里面的unfinished_tasks加1，get后需要task_done()才会减1

        self._client.close()
        self._closed = True
        time.sleep(1)
        print("SSHClient exit~~~~~~~")

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


if __name__ == '__main__':
    # 测试代码
    config = {
        "hostname": '10.0.0.179',
        "port": 22,
        "username": "root",  # 在目标主机上的用户身份，如果不写就和当前客户端用户同名
        # "password": "zjwjl2004",
        "ssh_pkey_path": "/var/www/nexora/uploads/2025/11/08/61ad5024445840f8a2a9e78de9948030"  # 指定成对的密钥的私钥文件
    }


    def get_remote_output(client: SSHClient):
        while not client.closed:
            data = client.output()
            if data is not None:
                print(f"data = {data}")


    with SSHClient(**config) as client:
        threading.Thread(target=get_remote_output, args=(client,)).start()
        time.sleep(1)
        for c in ('id', 'pwd', 'ls -l /dev\r'):
            client.send(f'{c}\r')
            time.sleep(1)
        time.sleep(2)
        client.send('exit\r')
        time.sleep(1)
